import time, re, tesserocr
import numpy as np
from selenium import webdriver
from io import BytesIO
from PIL import Image
from retrying import retry
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.chrome.service import Service

def preprocess(image):
    # 转为灰度图
    image = image.convert('L')
    # 数组化
    array = np.array(image)
    # 将灰度阈值大于50的置为白色 255，小于50的置为黑色 0
    array = np.where(array > 100, 255, 0)
    # Image.fromarray()函数将数组转为图像
    # array.astype('uint8')函数将数组的元素数据类型转换为uint8，以确保像素值在0到255的范围内
    image = Image.fromarray(array.astype('uint8'))
    return image

@retry(stop_max_attempt_number=10, retry_on_result=lambda x: x is False)
def login():
    browser.get('https://captcha7.scrape.center/')
    browser.find_element(By.CSS_SELECTOR, '.username input[type="text"]').send_keys('admin')
    browser.find_element(By.CSS_SELECTOR, '.password input[type="password"]').send_keys('admin')
    captcha = browser.find_element(By.CSS_SELECTOR, '#captcha')
    # 找到并截取验证码图片，转化为图片对象
    image = Image.open(BytesIO(captcha.screenshot_as_png))
    # 预处理验证码，去除噪声
    image = preprocess(image)
    # 识别验证码，得到识别结果
    captcha = tesserocr.image_to_text(image)
    # 去除识别结果中的一些非字母字符和数字字符
    captcha = re.sub('[^A-Za-z0-9]', '', captcha)
    browser.find_element(By.CSS_SELECTOR, '.captcha input[type="text"]').send_keys(captcha)
    browser.find_element(By.CSS_SELECTOR, '.login').click()

    try:
        WebDriverWait(browser, 10).until(EC.presence_of_element_located((By.XPATH, '//h2[contains(.,"登录成功")]')))
        time.sleep(10)
        browser.close()
        return True
    except TimeoutException:
        return False




if __name__ == '__main__':
    driver_path = 'D:\install\Python\Scripts\chromedriver.exe'
    browser = webdriver.Chrome(service=Service(driver_path))
    login()